package com.xd.springaiagent.agent;

import com.xd.springaiagent.agent.AgentState;
import lombok.Data;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

@Data
@Component
public abstract class BaseAgent {
    protected String name;
    protected String description;
    protected String systemPrompt;
    protected String nextStepPrompt;
    protected ChatClient chatClient;
    protected AgentState state = AgentState.IDLE;
    protected List<Message> msg = new ArrayList<>();

    protected int maxSteps = 10;
    protected int currentStep = 0;

    protected void addMessage(String role, String content) {
        switch (role) {
            case "user":
                msg.add(new UserMessage(content));
                break;
            case "assistant":
                msg.add(new AssistantMessage(content));
                break;
            case "system":
                msg.add(new SystemMessage(content));
                break;
            default:
                throw new IllegalArgumentException("Unsupported message role: " + role);
        }
    }

    public void run(String request, SseEmitter emitter) {
        if (state != AgentState.IDLE) {
            throw new RuntimeException("无法从状态 " + state.getState() + " 运行 Agent");
        }

        addMessage("user", request);
        setState(AgentState.RUNNING);

        try {
            while (currentStep < maxSteps && state != AgentState.FINISHED) {
                currentStep++;
                String stepResult = step();
                String message = String.format("{\"step\": %d, \"result\": %s}", currentStep, stepResult); // 格式化为 JSON
                try {
                    emitter.send(SseEmitter.event()
                            .name("step" + currentStep)
                            .data(message)
                            .id(String.valueOf(currentStep))); // 添加 ID 便于跟踪
                    emitter.send("\n\n"); // 结束数据块
                } catch (IOException e) {
                    setState(AgentState.ERROR);
                    emitter.send(SseEmitter.event()
                            .name("error")
                            .data("{\"error\": \"发送数据失败: " + e.getMessage() + "\"}")
                            .id("error"));
                    emitter.completeWithError(e);
                    return;
                }
            }
            emitter.complete();
        } catch (Exception e) {
            setState(AgentState.ERROR);
            try {
                emitter.send(SseEmitter.event()
                        .name("error")
                        .data("{\"error\": \"错误: " + e.getMessage() + "\"}")
                        .id("error"));
                emitter.completeWithError(e);
            } catch (IOException ioException) {
                ioException.printStackTrace();
            }
        } finally {
            cleanup();
        }
    }

    protected abstract String step();

    protected abstract void cleanup();
}